package com.zhny.test;

import org.apache.parquet.Strings;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.DoubleFunction;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.PairFunction;
import org.apache.spark.mllib.linalg.Vector;
import org.apache.spark.mllib.linalg.Vectors;
import org.apache.spark.mllib.regression.GeneralizedLinearModel;
import org.apache.spark.mllib.regression.LabeledPoint;
import org.apache.spark.mllib.regression.LinearRegressionModel;
import org.apache.spark.mllib.regression.LinearRegressionWithSGD;

import scala.Tuple2;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
//线性回归
public class LinearRegressionUtil {
    public static void exc(String dataFilePath, String resultFilePath) {
        SparkConf sparkConf = new SparkConf().setAppName("LinearRegression Algorithm").setMaster("local[*]");
        sparkConf.set("spark.driver.allowMultipleContexts", "true");
        JavaSparkContext sc = new JavaSparkContext(sparkConf);
        JavaRDD<String> data = null;

        if (Strings.isNullOrEmpty(dataFilePath)) {
            data = sc.textFile("src/main/resources/data/LinearRegressionData.txt");
        } else {
            data = sc.textFile(dataFilePath);
        }

        JavaRDD<LabeledPoint> parsedData = data.map(new Function<String, LabeledPoint>() {
            private static final long serialVersionUID = 1L;

            @Override
            public LabeledPoint call(String line) throws Exception {
                String[] parts = line.split(",");

                double y = Double.parseDouble(parts[0]) / 10000;

                String[] xnStr = parts[1].split(" ");
                double[] xn = new double[xnStr.length];

                for (int i = 0; i < xn.length; i++) {
                    xn[i] = Double.parseDouble(xnStr[i]) / 10000;
                }

                return new LabeledPoint(y, Vectors.dense(xn));
            }
        }).cache();

        int numIterations = 100;

        LinearRegressionModel model = LinearRegressionWithSGD.train(parsedData.rdd(), numIterations, 10);
//		print(parsedData, model);
        printMean(parsedData, model);

        FileWriter fos = null;
        try {
            fos = new FileWriter(new File(resultFilePath));
            fos.write(model.toPMML() + "\n");

            double[] d = new double[] {1000.0, 1000.0, 1000.0, 1000.0};
            Vector v = Vectors.dense(d);

            fos.write(String.valueOf(new Double(model.predict(v)).longValue()) + "\n");

            fos.flush();
            fos.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static void print(JavaRDD<LabeledPoint> parsedData, final GeneralizedLinearModel model) {
        parsedData.map(new Function<LabeledPoint, LabeledPoint>() {
            private static final long serialVersionUID = 1L;

            @Override
            public LabeledPoint call(LabeledPoint arg0) throws Exception {
                double prediction = model.predict(arg0.features());
                System.out.println(arg0.label()+"->"+prediction);
                return null;
            }
        }).collect();
    }

    public static void printMean(JavaRDD<LabeledPoint> parsedData, final GeneralizedLinearModel model) {
        JavaPairRDD<Double, Double> valuesAndPreds = parsedData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
            private static final long serialVersionUID = 1L;
            @Override
            public Tuple2<Double, Double> call(LabeledPoint point) throws Exception {
                double prediction = model.predict(point.features());

                return new Tuple2<>(point.label(), prediction);
            }
        });


        Double MSE = valuesAndPreds.mapToDouble(new DoubleFunction<Tuple2<Double, Double>>() {
            private static final long serialVersionUID = 1L;
            @Override
            public double call(Tuple2<Double, Double> t) throws Exception {
                return Math.pow(t._1() - t._2(), 2);
            }
        }).mean();

        System.out.println(model.getClass().getName() + " training Mean Squared Error = " + MSE);
    }
}
